import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime, timedelta

# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')


# 获取股票池数据（这里以上证50成分股为例）
def get_stock_pool():
    # 获取上证50成分股
    df = pro.index_weight(index_code='000016.SH', start_date='20200101', end_date='20201231')
    stock_pool = df['con_code'].unique().tolist()
    return stock_pool


# 获取股票日线数据
def get_daily_data(ts_code, start_date, end_date):
    df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
    df = df.sort_values('trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    return df


# 计算技术指标
def calculate_technical_indicators(df):
    # 计算均线
    df['ma5'] = df['close'].rolling(5).mean()
    df['ma20'] = df['close'].rolling(20).mean()

    # 计算MACD
    df['ema12'] = df['close'].ewm(span=12, adjust=False).mean()
    df['ema26'] = df['close'].ewm(span=26, adjust=False).mean()
    df['macd'] = df['ema12'] - df['ema26']
    df['signal'] = df['macd'].ewm(span=9, adjust=False).mean()
    df['hist'] = df['macd'] - df['signal']

    # 计算RSI
    delta = df['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    df['rsi'] = 100 - (100 / (1 + rs))

    return df.dropna()


# 生成交易信号
def generate_signals(df):
    signals = pd.DataFrame(index=df.index)
    signals['price'] = df['close']

    # 双均线策略信号
    signals['ma5'] = df['ma5']
    signals['ma20'] = df['ma20']
    signals['ma_signal'] = np.where(signals['ma5'] > signals['ma20'], 1, 0)

    # MACD策略信号
    signals['macd_hist'] = df['hist']
    signals['macd_signal'] = np.where(signals['macd_hist'] > 0, 1, 0)

    # RSI策略信号
    signals['rsi'] = df['rsi']
    signals['rsi_signal'] = np.where(signals['rsi'] < 30, 1,
                                     np.where(signals['rsi'] > 70, -1, 0))

    # 综合信号
    signals['combined_signal'] = signals['ma_signal'] + signals['macd_signal'] + signals['rsi_signal']

    # 生成买卖信号 (1:买, -1:卖, 0:持有)
    signals['positions'] = signals['combined_signal'].diff()

    return signals


# 选股与排序
def select_and_rank_stocks(stock_pool, date):
    ranked_stocks = []

    for stock in stock_pool:
        try:
            # 获取最近60个交易日数据
            end_date = datetime.strptime(date, '%Y%m%d')
            start_date = end_date - timedelta(days=120)
            df = get_daily_data(stock, start_date.strftime('%Y%m%d'), date)

            if len(df) < 60:  # 确保有足够数据
                continue

            df = calculate_technical_indicators(df)

            # 计算选股指标
            last_row = df.iloc[-1]

            # 动量指标 (20日收益率)
            momentum = (last_row['close'] / df.iloc[-20]['close'] - 1) * 100

            # 波动率 (20日标准差)
            volatility = df['close'].pct_change().rolling(20).std().iloc[-1]

            # 成交量变化
            volume_change = (last_row['vol'] / df['vol'].rolling(20).mean().iloc[-1] - 1) * 100

            # 综合评分 (可根据需要调整权重)
            score = momentum * 0.5 - volatility * 0.3 + volume_change * 0.2

            ranked_stocks.append({
                'ts_code': stock,
                'momentum': momentum,
                'volatility': volatility,
                'volume_change': volume_change,
                'score': score
            })

        except Exception as e:
            print(f"Error processing {stock}: {e}")
            continue

    # 按分数排序
    ranked_df = pd.DataFrame(ranked_stocks)
    ranked_df = ranked_df.sort_values('score', ascending=False)

    return ranked_df


# 止损止盈策略
def apply_stop_loss_take_profit(positions, prices, stop_loss_pct=0.05, take_profit_pct=0.1):
    """
    应用止损止盈策略
    :param positions: 原始仓位信号
    :param prices: 价格序列
    :param stop_loss_pct: 止损百分比
    :param take_profit_pct: 止盈百分比
    :return: 调整后的仓位信号
    """
    adjusted_positions = positions.copy()
    holding = False
    entry_price = 0

    for i in range(1, len(positions)):
        if not holding and positions[i] == 1:  # 新开仓
            holding = True
            entry_price = prices[i]
            adjusted_positions[i] = 1
        elif holding:
            current_price = prices[i]
            # 计算盈亏比例
            pct_change = (current_price - entry_price) / entry_price

            # 检查止损或止盈条件
            if pct_change <= -stop_loss_pct or pct_change >= take_profit_pct:  # 这里修正了变量名
                adjusted_positions[i] = -1
                holding = False
            else:
                adjusted_positions[i] = 0  # 保持仓位
        else:
            adjusted_positions[i] = 0

    return adjusted_positions


# 回测引擎
def backtest(signals, initial_capital=1000000):
    portfolio = pd.DataFrame(index=signals.index)
    portfolio['price'] = signals['price']
    portfolio['positions'] = signals['positions']

    # 计算持仓
    portfolio['holdings'] = portfolio['positions'].cumsum()

    # 计算现金变化
    portfolio['cash'] = initial_capital - (portfolio['positions'] * portfolio['price']).cumsum()

    # 计算总资产
    portfolio['total'] = portfolio['cash'] + (portfolio['holdings'] * portfolio['price'])

    # 计算收益率
    portfolio['returns'] = portfolio['total'].pct_change()

    return portfolio


# 可视化结果
def visualize_backtest(portfolio, signals):
    plt.figure(figsize=(14, 10))

    # 价格和均线
    ax1 = plt.subplot(311)
    portfolio['price'].plot(ax=ax1, label='Price')
    signals['ma5'].plot(ax=ax1, label='5-day MA')
    signals['ma20'].plot(ax=ax1, label='20-day MA')

    # 买卖信号
    buy_signals = signals[signals['positions'] == 1]
    sell_signals = signals[signals['positions'] == -1]
    ax1.plot(buy_signals.index, signals.loc[buy_signals.index, 'price'], '^', markersize=10, color='g', label='Buy')
    ax1.plot(sell_signals.index, signals.loc[sell_signals.index, 'price'], 'v', markersize=10, color='r', label='Sell')

    ax1.set_title('Price and Trading Signals')
    ax1.legend()

    # 持仓情况
    ax2 = plt.subplot(312)
    portfolio['holdings'].plot(ax=ax2, label='Holdings')
    ax2.set_title('Position Holdings')
    ax2.legend()

    # 资产变化
    ax3 = plt.subplot(313)
    portfolio['total'].plot(ax=ax3, label='Total Value')
    ax3.set_title('Portfolio Value')
    ax3.legend()

    plt.tight_layout()
    plt.show()

    # 显示关键指标
    total_return = (portfolio['total'].iloc[-1] / initial_capital - 1) * 100
    max_drawdown = (portfolio['total'].cummax() - portfolio['total']).max() / portfolio['total'].cummax().max() * 100
    sharpe_ratio = portfolio['returns'].mean() / portfolio['returns'].std() * np.sqrt(252)

    print(f"总收益率: {total_return:.2f}%")
    print(f"最大回撤: {max_drawdown:.2f}%")
    print(f"夏普比率: {sharpe_ratio:.2f}")


# 主程序
if __name__ == "__main__":
    # 参数设置
    start_date = '20200101'
    end_date = '20201231'
    initial_capital = 1000000

    # 1. 获取股票池
    stock_pool = get_stock_pool()
    print(f"股票池数量: {len(stock_pool)}")

    # 2. 选股与排序 (以2020年6月1日为例)
    ranked_stocks = select_and_rank_stocks(stock_pool, '20200601')
    print("排名前10的股票:")
    print(ranked_stocks.head(10))

    # 3. 选择排名第一的股票进行回测
    selected_stock = ranked_stocks.iloc[0]['ts_code']
    print(f"\n选择股票: {selected_stock} 进行回测")

    # 4. 获取数据并计算指标
    df = get_daily_data(selected_stock, start_date, end_date)
    df = calculate_technical_indicators(df)

    # 5. 生成交易信号
    signals = generate_signals(df)

    # 6. 应用止损止盈策略
    signals['positions'] = apply_stop_loss_take_profit(signals['positions'], signals['price'])

    # 7. 回测
    portfolio = backtest(signals, initial_capital)

    # 8. 可视化结果
    visualize_backtest(portfolio, signals)